import cv2
import os

from TrainFace import trainFace


trainFace()
# 加载训练集
recognizer = cv2.face.LBPHFaceRecognizer_create() # LBPH（局部二值模式直方图）
recognizer.read("./data/train.yml")

# 将id与姓名进行绑定(后面可以检索出该用户的名称)
def name():
    """
    解析 ./train/ 目录下的人脸图像文件名，并建立 ID 与姓名的对应关系。
    文件命名格式应为：{id}.{name}.jpg，例如：1.Alice.jpg
    """

    names = {}  # 创建一个空字典，用于存储 ID 和姓名的映射关系

    # 获取 ./train/ 目录下所有文件的完整路径
    filePaths = [os.path.join('./train/', f) for f in os.listdir('./train/')]

    # 遍历每个文件路径
    for fpath in filePaths:
        # 获取文件名（去掉路径），并按照 '.' 进行拆分
        filename = os.path.split(fpath)[1]  # 提取文件名，例如 "1.Alice.jpg"
        id = str(filename.split(".")[0])  # 获取 ID 部分，例如 "1"
        name = str(filename.split(".")[1])  # 获取姓名部分，例如 "Alice"

        # 将 ID 和姓名存入字典
        names[id] = name

    return names  # 返回 ID 和姓名的映射字典


#检测人脸
def faceDetection(frame):
    """
    使用 OpenCV 进行人脸检测，并在图像上绘制识别结果。

    主要流程：
    1. 加载人脸检测分类器（Haar 级联分类器）。
    2. 将输入的彩色图像转换为灰度图，提高检测精度。
    3. 进行人脸检测，获取所有人脸的位置信息。
    4. 使用训练好的识别模型对检测到的人脸进行身份识别。
    5. 在检测到的人脸区域绘制矩形框，并标注姓名或“unknown”。
    6. 显示处理后的图像。

    参数：
    frame : numpy.ndarray
        摄像头捕获的当前帧（BGR 彩色图像）。
    """

    # 加载 OpenCV 的 Haar 级联分类器进行人脸检测
    classifier = cv2.CascadeClassifier("./haarcascades/haarcascade_frontalface_alt2.xml")

    # 将彩色图像转换为灰度图，减少计算量，提高检测精度
    gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 进行人脸检测
    faces = classifier.detectMultiScale(
        gray_img,  # 输入灰度图
        scaleFactor=1.01,  # 图像缩放比例，每次检测时缩小图像尺寸（提高检测效率）
        minNeighbors=5,  # 最少相邻检测到的矩形数，值越高误检越少，但可能漏检
        flags=0,  # 一般设置为 0，OpenCV 可能会使用默认参数
        minSize=(100, 100),  # 识别人脸的最小尺寸（小于该尺寸的人脸忽略）
        maxSize=(500, 500)  # 识别人脸的最大尺寸（大于该尺寸的人脸忽略）
    )

    # 获取 id 和姓名的映射关系
    names = name()

    # 遍历检测到的所有人脸
    for face in faces:
        (x, y, w, h) = face  # 获取人脸的坐标 (x, y) 以及宽度 w、高度 h

        # 在检测到的人脸区域绘制矩形框
        cv2.rectangle(frame, (x, y), (x + w, y + h), color=(255, 0, 0), thickness=2)

        # 使用训练好的 LBPH 人脸识别器对检测到的人脸进行身份预测
        id, confidence = recognizer.predict(gray_img[y:y + h, x:x + w])

        # confidence 置信度：值越低，代表识别越准确
        if confidence < 80:  # 设置 80 作为阈值，小于 80 认为是匹配的
            cv2.putText(
                frame,  # 目标图像
                names.get(str(id), "unknown"),  # 获取 ID 对应的姓名（若找不到则显示 "unknown"）
                (x, y),  # 文字坐标（左上角）
                cv2.FONT_HERSHEY_SIMPLEX,  # 字体样式
                1.2,  # 字体大小
                (255, 0, 0),  # 文字颜色（蓝色）
                2  # 线条粗细
            )
        else:
            # 若置信度高于 80，则标记为 "unknown"
            cv2.putText(
                frame,
                "unknown",
                (x, y),
                cv2.FONT_HERSHEY_SIMPLEX,
                1.2,
                (255, 0, 0),
                2
            )

    # 显示带有识别结果的图像窗口
    cv2.imshow("frame", frame)


if __name__ == '__main__':  # 只有当此脚本作为主程序运行时，才会执行下面的代码
    # 打开摄像头，参数 0 代表默认摄像头
    cap = cv2.VideoCapture(0)

    while True:  # 进入一个循环，持续获取摄像头画面
        fig, frame = cap.read()  # 读取摄像头捕获的帧，fig 表示是否成功读取，frame 是实际图像数据

        if fig is None:  # 如果未成功捕获到图像，打印错误信息并退出循环
            print("Capture No Found")
            break

        frame = cv2.flip(frame, 1)  # 水平翻转帧，修正摄像头默认的镜像问题

        if cv2.waitKey(1) == ord(" "):  # 如果按下空格键（" "），则退出循环
            break

        faceDetection(frame)  # 调用 faceDetection() 处理当前帧并检测人脸

    # 释放摄像头资源
    cap.release()

    # 关闭所有 OpenCV 窗口
    cv2.destroyAllWindows()
